import os
import numpy as np
from scipy.stats import norm
from PIL import Image

def scale(data):
    """将数据线性缩放到0-255的整数范围（灰度均衡）"""
    if not data:
        return []
    min_val = min(data)
    max_val = max(data)
    if max_val == min_val:
        return [0] * len(data)  # 避免除以0
    # 线性缩放公式：(x - min) * 255 / (max - min)
    scaled = [(x - min_val) * 255.0 / (max_val - min_val) for x in data]
    return [int(round(x)) for x in scaled]  # 四舍五入为整数


def calcEM(data):
    """
    用EM算法拟合两个高斯分量的混合模型（GMM）
    参数：data（一维列表）：图像通道的像素数据
    返回：parameter（列表）：模型参数，顺序为[weight1, weight2, mu1, sigma1, mu2, sigma2]
          - weight1/2：两个分量的权重（之和≈1）
          - mu1/2：两个分量的均值（对应像素值的两个峰值）
          - sigma1/2：两个分量的标准差（对应像素值的分散程度）
    """
    N = len(data)
    # 初始化参数（经验值：权重各0.5，均值取最小值/最大值，标准差为1）
    weight1, weight2 = 0.5, 0.5
    mu1, mu2 = min(data), max(data)
    sigma1, sigma2 = 1.0, 1.0

    # 初始化后验概率（每个样本属于两个分量的概率）
    gamma1 = [0.0] * N
    gamma2 = [0.0] * N

    # 迭代停止条件：最大迭代100次或参数变化小于阈值
    max_iter = 100
    tolerance = 1e-6

    for iter in range(max_iter):
        # ------------------- E步：计算后验概率 -------------------
        for i in range(N):
            x = data[i]
            # 计算两个分量的似然度（概率密度）
            pdf1 = norm.pdf(x, loc=mu1, scale=sigma1)  # 分量1的概率密度
            pdf2 = norm.pdf(x, loc=mu2, scale=sigma2)  # 分量2的概率密度
            # 归一化（避免分母为0）
            denominator = weight1 * pdf1 + weight2 * pdf2
            if denominator == 0:
                gamma1[i], gamma2[i] = 0.0, 0.0
            else:
                gamma1[i] = (weight1 * pdf1) / denominator  # 属于分量1的概率
                gamma2[i] = (weight2 * pdf2) / denominator  # 属于分量2的概率

        # ------------------- M步：更新参数 -------------------
        # 有效样本数（加权和）
        N1 = sum(gamma1)
        N2 = sum(gamma2)
        # 更新权重（有效样本数/总样本数）
        new_weight1 = N1 / N if N1 > 0 else weight1
        new_weight2 = N2 / N if N2 > 0 else weight2
        # 更新均值（加权平均）
        new_mu1 = sum(gamma1[i] * data[i] for i in range(N)) / N1 if N1 > 0 else mu1
        new_mu2 = sum(gamma2[i] * data[i] for i in range(N)) / N2 if N2 > 0 else mu2
        # 更新标准差（加权标准差）
        new_sigma1 = np.sqrt(sum(gamma1[i] * (data[i] - new_mu1) ** 2 for i in range(N)) / N1) if N1 > 0 else sigma1
        new_sigma2 = np.sqrt(sum(gamma2[i] * (data[i] - new_mu2) ** 2 for i in range(N)) / N2) if N2 > 0 else sigma2

        # ------------------- 检查收敛 -------------------
        if (abs(new_weight1 - weight1) < tolerance and
                abs(new_weight2 - weight2) < tolerance and
                abs(new_mu1 - mu1) < tolerance and
                abs(new_sigma1 - sigma1) < tolerance and
                abs(new_mu2 - mu2) < tolerance and
                abs(new_sigma2 - sigma2) < tolerance):
            break

        # 更新参数
        weight1, weight2 = new_weight1, new_weight2
        mu1, mu2 = new_mu1, new_mu2
        sigma1, sigma2 = new_sigma1, new_sigma2

    return [weight1, weight2, mu1, sigma1, mu2, sigma2]


def composite(band, parameter):
    """
    生成两个高斯分量的复合通道（概率密度图）
    参数：
        band（列表）：图像通道的像素数据（一维）
        parameter（列表）：GMM模型参数（顺序与`calcEM`返回一致）
    返回：
        list：两个缩放后的复合通道数据（0-255整数，对应两个分量的概率密度图）
    """
    # 从参数中提取两个高斯分量的信息（顺序与`calcEM`返回一致）
    weight1, weight2, mu1, sigma1, mu2, sigma2 = parameter

    # 计算每个像素的加权概率密度（分量1和分量2）
    p1 = []  # 分量1的加权概率密度（weight1 * pdf1）
    p2 = []  # 分量2的加权概率密度（weight2 * pdf2）
    for pixel in band:
        pdf1 = norm.pdf(pixel, loc=mu1, scale=sigma1)  # 分量1的概率密度
        pdf2 = norm.pdf(pixel, loc=mu2, scale=sigma2)  # 分量2的概率密度
        p1.append(weight1 * pdf1)
        p2.append(weight2 * pdf2)

    # 将概率密度值缩放到0-255（灰度均衡）
    p1_scaled = scale(p1)
    p2_scaled = scale(p2)

    return [p1_scaled, p2_scaled]


if __name__ == "__main__":
    # ------------------- 1. 读取图像 -------------------
    img_path = os.path.join(".", "Pic", "test.bmp")  # 图像路径（跨平台）
    try:
        im = Image.open(img_path)
    except IOError:
        print(f"错误：无法打开图像文件「{img_path}」")
        exit()

    # 打印图像信息
    print(f"图像格式：{im.format}")
    print(f"图像尺寸：{im.size}（宽×高）")
    print(f"图像模式：{im.mode}（如RGB、L）")

    # ------------------- 2. 提取通道数据 -------------------
    if im.mode in ["RGB", "RGBA", "CMYK"]:
        # 分割通道（如RGB图像取第一个通道：R通道）
        channels = im.split()
        if not channels:
            print("错误：无法分割通道")
            exit()
        channel = channels[0]  # 取第一个通道（可根据需求调整）
    else:
        # 灰度图像直接使用
        channel = im

    # 获取通道的像素数据（一维列表，长度=宽×高）
    data = list(channel.getdata())
    print(f"通道像素数量：{len(data)}")

    # ------------------- 3. 用GMM拟合数据 -------------------
    print("\n正在用GMM拟合数据...")
    parameter = calcEM(data)
    print(f"GMM模型参数：{parameter}")
    print(f"  - 分量1：权重={parameter[0]:.2f}，均值={parameter[2]:.2f}，标准差={parameter[3]:.2f}")
    print(f"  - 分量2：权重={parameter[1]:.2f}，均值={parameter[4]:.2f}，标准差={parameter[5]:.2f}")

    # ------------------- 4. 生成复合通道 -------------------
    print("\n正在生成复合通道...")
    composite_channels = composite(data, parameter)
    p1_scaled, p2_scaled = composite_channels  # 两个复合通道（0-255整数）

    # ------------------- 5. 保存复合通道图像 -------------------
    # 生成分量1的概率密度图（灰度图像）
    img1 = Image.new("L", channel.size)  # "L"模式：8位灰度图像
    img1.putdata(p1_scaled)
    img1.save(os.path.join(".", "Pic", "composite_channel1.bmp"))
    print("已保存分量1的概率密度图：composite_channel1.bmp")

    # 生成分量2的概率密度图（灰度图像）
    img2 = Image.new("L", channel.size)
    img2.putdata(p2_scaled)
    img2.save(os.path.join(".", "Pic", "composite_channel2.bmp"))
    print("已保存分量2的概率密度图：composite_channel2.bmp")

    print("\n处理完成！")